package common;

import java.util.ArrayList;
import java.util.TreeMap;

public class CorpusProcessor {
	private TreeMap<String, TreeMap<String, Integer>> wordCounts;
	private TreeMap<String, Integer> posCounts;
	private TreeMap<Bigram, Integer> bigramCounts;
	
	public CorpusProcessor(){
		wordCounts = new TreeMap<String, TreeMap<String, Integer>>();
		posCounts = new TreeMap<String, Integer>();
		bigramCounts = new TreeMap<Bigram, Integer>();
	}
	
	public void process(ArrayList<String> words, ArrayList<String> pos){
		TreeMap<String, Integer> tagCount;
		for(int i = 0; i < words.size(); i++){
			// Extract distinct words and frequencies
			String thisWord = words.get(i);
			String thisPos = pos.get(i);
			if(wordCounts.containsKey(thisWord)){
				tagCount = wordCounts.get(thisWord);
				int count1 = tagCount.containsKey(thisPos) ? tagCount.get(thisPos) : 0;
				tagCount.put(thisPos, count1 + 1);
				wordCounts.put(thisWord, tagCount);
			}else{
				tagCount = new TreeMap<String, Integer>();
				tagCount.put(thisPos, 1);
				wordCounts.put(thisWord, tagCount);
			}

			// Extract distinct pos and frequencies
			int count2 = posCounts.containsKey(thisPos) ? posCounts.get(thisPos) : 0;
			posCounts.put(thisPos, count2 + 1);
			posCounts.put("<s>", 1);
			
			// Extract distinct pos bigrams and frequencies
			Bigram b;
			if(i != 0){
				b = new Bigram(thisPos, pos.get(i-1));
			}else{
				b = new Bigram(thisPos, "<s>");
			}
			int count3 = bigramCounts.containsKey(b) ? bigramCounts.get(b) : 0;
			bigramCounts.put(b, count3 + 1);
		}
	}
	
	public TreeMap<String, TreeMap<String, Integer> > getWordCounts(){
		return wordCounts;
	}
	
	public TreeMap<String, Integer> getPosCounts(){
		return posCounts;
	}
	
	public TreeMap<Bigram, Integer> getBigramCounts(){
		return bigramCounts;
	}
}
